"""
Case Type   : DATAVEC模块
Case Name   : bit字段创建ivfflat 汉明距离索引，索引选项lists参数覆盖，lists参数不同，索引查询准确度验证
Create At   : 2025/4/18
Owner       : @peilinqian
Description :
    step1:创建3张表;
    step2:创建自定义插入函数;
    step3:表插入相同数据;
    step4:创建索引，指定lists，小于默认值100;
    step5:创建索引，不指定lists;
    step6:创建索引，指定lists，大于默认值100;
    step7: ivfflat_probes设置较小,不同参数索引查询召回率
Expect      :
    step1:创建3张表; expect:成功
    step2:创建自定义插入函数; expect:成功
    step3:表插入相同数据; expect:成功
    step4:创建索引，指定lists，小于默认值100; expect:成功
    step5:创建索引，不指定lists; expect:成功
    step6:创建索引，指定lists，大于默认值100; expect:成功
    step7:ivfflat_probes设置较小,不同参数索引查询召回率; expect:list值与ivfflat_probes相差越大，召回率越低
History     :
"""
import os
import unittest

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node


class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        self.log = Logger()
        self.log.info(f'----{os.path.basename(__file__)}:start----')
        self.sh = CommonSH('PrimaryDbUser')
        self.pri_dbuser = Node(node='PrimaryDbUser')
        self.sta_1_sh = CommonSH('Standby1DbUser')
        self.constant = Constant()
        self.com = Common()
        self.comsql = CommonSQL('PrimaryDbUser')
        self.tb_name = 't_datavec_index_case0144'
        self.fun_insert = 'f_datavec_index_case0144'

    def test_main(self):
        step_txt = '----step1:创建3张表; expect:成功----'
        self.log.info(step_txt)
        create_tb = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                    f"{self.tb_name}_3;" \
                    f"create table {self.tb_name}_1 " \
                    f"(id bigserial primary key, val bit(200));" \
                    f"create table {self.tb_name}_2 " \
                    f"(id bigserial primary key, val bit(200));" \
                    f"create table {self.tb_name}_3 " \
                    f"(id bigserial primary key, val bit(200));"
        result = self.sh.execut_db_sql(create_tb)
        self.log.info(result)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result,
                      '执行失败:' + step_txt)

        step_txt = '----step2:创建自定义插入函数; expect:成功----'
        self.log.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        self.log.info(result)
        self.assertIn(self.fun_insert, result, '执行失败:' + step_txt)

        step_txt = '----step3:表插入相同数据; expect:成功----'
        self.log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(200,5000,'{self.tb_name}_1'," \
                     f"'val','bit');"
        result = self.sh.execut_db_sql(insert_sql)
        self.log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        result = self.sh.execut_db_sql(insert_sql)
        self.log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_2 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        self.log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_3 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        self.log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step4:创建索引，指定lists，小于默认值100; expect:成功----'
        self.log.info(step_txt)
        create_index = f"set maintenance_work_mem='100MB';" \
                       f"create index idx_{self.tb_name}_1 on " \
                       f"{self.tb_name}_1 using ivfflat" \
                       f"(val bit_hamming_ops) with (lists=50);"
        result = self.sh.execut_db_sql(create_index)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step5:创建索引，不指定lists; expect:成功----'
        self.log.info(step_txt)
        create_index = f"set maintenance_work_mem='100MB';" \
                       f"create index idx_{self.tb_name}_2 on " \
                       f"{self.tb_name}_2 using ivfflat(val bit_hamming_ops);"
        result = self.sh.execut_db_sql(create_index)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step6:创建索引，指定lists，大于默认值100; expect:成功----'
        self.log.info(step_txt)
        create_index = f"set maintenance_work_mem='100MB';" \
                       f"create index idx_{self.tb_name}_3 on " \
                       f"{self.tb_name}_3 using ivfflat" \
                       f"(val bit_hamming_ops) with(lists=300);"
        result = self.sh.execut_db_sql(create_index)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step7: ivfflat_probes设置较小,不同参数索引查询召回率; ' \
                   'expect:list值与ivfflat_probes相差越大，召回率越低----'
        self.log.info(step_txt)
        bit_value = '111101000011111001101011110101011101000111011011011001' \
                    '110010111000110011111011001010000011000011111101100001' \
                    '111110001110111101010000011101001101111001010001000110' \
                    '10000110011100100110010101101100110000'
        sql = f"select id, val <~> '{bit_value}' " \
              f"from {self.tb_name}_1 order by val <~> " \
              f"'{bit_value}' limit 20;"
        result1 = self.comsql.vector_query_recall(
            sql, 'set ivfflat_probes=50;', is_detail=False)
        sql = f"select id, val <~> '{bit_value}' " \
              f"from {self.tb_name}_2 order by val <~> " \
              f"'{bit_value}' limit 20;"
        result2 = self.comsql.vector_query_recall(
            sql, 'set ivfflat_probes=50;', is_detail=False)
        sql = f"select id, val <~> '{bit_value}' " \
              f"from {self.tb_name}_3 order by val <~> " \
              f"'{bit_value}' limit 20;"
        result3 = self.comsql.vector_query_recall(
            sql, 'set ivfflat_probes=50;', is_detail=False)
        self.assertGreaterEqual(result1, result2, '执行失败:' + step_txt)
        self.assertGreaterEqual(result1, result3, '执行失败:' + step_txt)

    def tearDown(self):
        step_txt = '----this is teardown----'
        self.log.info(step_txt)
        step1_txt = '----teardown:清理测试数据; expect:成功----'
        self.log.info(step1_txt)
        drop_cmd = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                   f"{self.tb_name}_3;"
        drop_result1 = self.sh.execut_db_sql(drop_cmd)
        self.log.info(drop_result1)
        drop_cmd = f"drop function if exists {self.fun_insert};"
        drop_result2 = self.sh.execut_db_sql(drop_cmd)
        self.log.info(drop_result2)

        self.log.info(f'----{os.path.basename(__file__)}:end----')
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, drop_result1,
                      '执行失败:' + step_txt + step1_txt)
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, drop_result2,
                      '执行失败:' + step_txt + step1_txt)
